@zyongye
I've cleaned up the old PR, could you help review this one?

💡 Codex Review

Keep DeepGEMM requirement for SM120 FP4 indexer path

This helper now disables the DeepGEMM requirement for every SM120 run, but the FP4 indexer cache path still depends on DeepGEMM kernels (fp8_fp4_*) because the new SM120 fallback only handles q_scale is None (FP8 Q). With use_fp4_cache=True on SM120 and no DeepGEMM installed, construction succeeds and the first prefill/decode call fails at runtime with the DeepGEMM _missing() error instead of being rejected up front.

Propagate weight_name_filter to fast safetensor loaders

The new pre-load weight_name_filter is only wired into safetensors_weights_iterator; this branch still loads all tensors for fastsafetensors (and similarly other non-default safetensor iterators), so skipped tensors are still materialized. For DeepSeek V4 this defeats the intended early skip of MTP weights and can reintroduce high transient memory use/OOM when these load formats are enabled.

@GanyX19 — both reproduced and verified on our GB10 (sm_121, 2-node TP=2), and all three are now landed on this branch (codex/ds4-sm120-min-enable, commits 15434cb64f + b43470e871, co-authored to you). Excellent, careful report — thank you.

#1 — MQA-logits constexpr leak. Confirmed: the three kernels declare num_q/seq_len_kv, num_rows/logits_width, stride_lm as tl.constexpr (9 lines), and grepping their uses they only ever appear in boundary masks (offs < num_rows) and address arithmetic (* stride_lm) — never a tl.arange/range extent — so making them plain runtime args is safe, exactly as you said. stride_ln (=1) and the tile/shape constexprs stay. This matches host-freeze behaviour we'd independently hit under sustained load on unified memory.

A nice corroboration fell out of our A/B: in a controlled run the pre-fix (constexpr) build logged _tf32_hc_prenorm_gemm_kernel JIT-compiling during inference and then died with EngineDeadError under the benchmark — i.e. the recompile pathology biting live. The fixed build ran the same load clean.

#3 — tf32 prenorm gemm. Took this too: same constexpr → runtime for stride_outs/stride_sqs (address-arithmetic only, no hint — store path, as you noted).

#2 — einsum alignment hint. Confirmed our branch already keeps a_stride_group/a_scale_stride_group runtime (per the do-not-specialize comment) but without a tl.multiple_of. Added the hint — and applied it to both the BF16 einsum kernel and its fused-quant sibling _deepseek_v4_sm12x_fp8_einsum_quant_kernel (same matmul, same A strides). Left a_scale_stride_hidden and the tf32 store strides un-hinted, per your reasoning about ÷16-provability.

Validation (GB10, DSv4-Flash MTP-2): correctness clean with the hints in — GSM8K-100 = 0.99, needle-recall 16/16 — so the multiple_of doesn't corrupt on our shapes. On decode the change is throughput-neutral at our benchmark context sizes (≤32K), and we didn't observe the −24%/recovery there. That lines up with your bisection being at 256K single-stream where the einsum is a much larger fraction of the step — it should recover exactly as you measured at long context; we just don't exercise that regime in the standard bench. The host-memory-leak → freeze fix is the headline win for us, and it's decode-neutral, so it's a clean take regardless.

Really appreciate the depth here — the mask/arithmetic-vs-extent distinction and the ÷16 safety analysis made these trivial to verify. Glad to have you running the branch on GB10.

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jasl.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Both fixes confirmed in production on our 4-node TP=4/EP=4 GB10 rig — thank you. And one remaining rough edge (long-context prefill latency) I'd love your read on.

9bccfd1 (persistent_topk) — wedge resolved. We've now run 11h+ at max_model_len=1048576 (full 1M) with zero wedges. The 1M-width cudagraph warmup — the exact spot that hard-failed on the oversubscribe
pre-fix — captures clean, 0 oversubscribe. And a pleasant surprise: on the current build the C128A workspace self-sizing shrank the per-context sparse-MLA workspaces enough that our KV pool is larger at 1M
(4.63M tokens) than it was at 512K (2.91M) — so we got the full context window back with more concurrency headroom, not less.

15434cb (MQA-logits constexpr leak) — the memory-creep headline win, confirmed. GanyX19's diagnosis was exactly our long-standing "unexplained ~2.4–3.5 MB/req creep to freeze." Post-fix, measured two ways:
the Triton compile cache now grows ~1.9 MB/hour (vs the pre-fix ~5–6 GB/hour — a ~3000× cut), and host avail holds flat (~16.8 GB across 8h) instead of marching to the earlyoom floor. The ~17h-restart
treadmill is over. (b43470e's tf32-prenorm runtime strides + einsum multiple_of are in too.)

The rough edge — long-context prefill monopolizes the engine (~1–2 min). Now that 1M is enabled, a single large prefill stalls everything. A representative window (2 requests running, ~1M-scale context):
prompt 0 tok/s, gen ~0 tok/s, Running 2, Waiting 0→2→4 (~1m50s)
prompt 11,609 tok/s, gen 51, Running 3, Waiting 1 (prefill lands, recovers)
For ~110s the engine emits ~0 tokens while grinding the prefill, the queue backs up to 4 waiting, then it completes in a burst and recovers. Not a wedge — clean recovery — but a real latency/throughput hit
under concurrency. Our read is it's the streaming-radix histogram_streaming_topk path (the re-stream-from-global-each-round you flagged) dominating the step at large context, compounded by decode being
interconnect-latency-bound here (NCCL-spin-wait, so we can't hide it behind other work). We tried chunked-prefill interleaving (smaller max_num_batched_tokens) earlier and it barely helped — even a 26K
prefill blocked a small request ~10s; the chunks are too heavy to interleave on this fabric.

Questions for you:

1. Is the streaming-radix long-context cost something you expect to optimize (e.g. a single-pass or cached-prefix variant), or is re-streaming inherent to fitting GB10's ~99 KB smem?
2. Any scheduling lever you'd recommend to bound the monopolization — long_prefill_token_threshold, max_num_partial_prefills, or something else — that actually interleaves on this hardware?
3. Given the einsum multiple_of hint recovers ~24% at long context, is more of the prefill step similarly recoverable, or is the top-k the hard floor?

Happy to run a controlled prefill-latency-vs-context sweep on the 4-node rig if that'd help you characterize it — we're a real 1M/EP=4 production case.

@wingcomm — this is fantastic to hear, thank you for the production confirmation. 11h+ at full 1M with zero wedges, the 1M cudagraph warmup capturing clean, and the ~5–6 GB/h → ~1.9 MB/h memory-creep cut (~3000×, "the 17h-restart treadmill is over") are exactly the outcomes we were hoping for — and the KV-pool bonus at 1M (4.63M vs 2.91M @512k from the C128A workspace self-sizing) is a nice surprise. Really appreciate you closing the loop with hard numbers on a real EP=4 rig.

On the long-context prefill monopolization — good diagnosis, and your read is mostly right. Answers:

1. Is the streaming-radix cost inherent to 99 KB smem? — No, it's a conservative floor, and there's real headroom.
The histogram_streaming_topk path re-streams the row from global on each of its 4 MSD-radix rounds — that's ~4× full-row reads at 1M, which is the cost you're seeing. We shipped it that way for exactness under the 99 KB cap (the bounded-buffer histogram_256 path under-selects at dense pivots — >4096 ties in a 256-bin threshold bin — which is why it's gated to ≤32 KB-context). But it doesn't have to re-read the whole row every round:

• After round 1 you already know the top 8-bit bin; rounds 2–4 only need the elements in that bin (~N/256), which for 1M is ~4096 — small enough to buffer in smem. So a "narrow-after-round-1" variant is ~1 full read + 3 passes over a few-K buffer, not 4 full reads.
• Even simpler: our own histogram_2048 (2048-bin, single-pass + bounded buffer) fits 99 KB and at 1M has only ~488 elements per threshold bin — well under the 4096 buffer, so it wouldn't overflow. It's currently capped at HIST2048_THRESHOLD=8192 conservatively; extending a finer single-pass histogram to the long-context force_noncooperative path (with the 4-round re-stream kept only as the rare genuinely-dense fallback) should recover most of the cost.

So: the re-stream is exact-and-safe, not minimal. I think it's optimizable and I'd like to prototype the finer-histogram/narrow variant — your sweep offer is exactly what would guide it (see below).

2. Scheduling lever — the reason chunked-prefill barely helped is structural, and it bounds the fix.
The indexer top-k is over the accumulated context, not the query chunk. So splitting a 1M prefill into chunks doesn't split the top-k work — each late-context chunk still runs a ~full-context top-k (the ~seconds you measured), it just reschedules it. Smaller max_num_batched_tokens interleaves finer but every late chunk is still heavy, so you hit a floor. The levers that do help are the partial-prefill ones, not the batched-token size: set max_num_partial_prefills > 1 and a long_prefill_token_threshold below your prompt size — that classifies the 1M prefill as "long" and lets short prompts jump ahead between its chunks (max_long_partial_prefills caps how many long ones run concurrently). That bounds the starvation (your Waiting 0→4 backing up) even though it can't shrink the long prefill's own cost. It's a latency-for-small-requests win, not a throughput win — the throughput fix is making the top-k cheaper (#1).

3. Is more of the prefill recoverable like the einsum? — No, different kind of floor.
The einsum win was a Triton alignment recovery (tl.multiple_of) — narrowed loads, free to restore. The top-k cost is algorithmic/bandwidth (the re-stream), so an alignment hint won't touch it; the recoverable part there is #1 (finer histogram). It's worth auditing the other long-context prefill kernels (sparse-MLA prefill, the MoE GEMMs) for their own alignment/warmup recoveries, but the top-k is the algorithmic hard part at 1M — and, as you note, decode being NCCL-spin-wait-bound on this fabric means you can't hide the prefill behind other work, which compounds it independently of the kernel.

On the sweep — yes please. A controlled prefill-latency-vs-context sweep on your 4-node 1M/EP=4 rig would be genuinely useful: it'd confirm whether the top-k actually dominates the step (vs the sparse-MLA prefill / MoE) at 1M, and give me a target curve to validate a finer-histogram prototype against. If you can capture per-context-length step time (say 128K / 256K / 512K / 1M, single large prefill, idle otherwise) plus a py-spy/nsys sample of where the ~110s goes, that pins it. If the top-k is the bulk, I'll take a run at the finer-histogram variant and send you a patch to try. We don't have a 4-node/1M setup here (our GB10 validation is 2-node), so your rig is the real characterization platform.

@wingcomm — didn't wait for the sweep: took a run at the finer-histogram top-k and it's landed on this branch (codex/ds4-sm120-min-enable, 9f18be7630), ready for you to try.

What it is. The 4-round streaming radix re-read the full row 4× (once per MSD-radix round). The new path (histogram_2048_topk<TopK, CacheBins=false> in persistent_topk.cuh) does the finer-histogram approach we discussed: Phase 1 builds the 2048-bin histogram, Phase 2 re-reads once and collects, Phase 3 exactly refines the buffered pivot bin. 2 full-row reads instead of 4. The finer bins keep the pivot bin small (~seq/2048 ≈ 488 at 1M, well under the 3708 tie buffer), so it stays single-CTA / 99 KB-smem-safe. If the pivot bin does overflow (fp16 bins clustering — e.g. thousands of exactly-equal logits), it falls back to the streaming radix, so it's provably exact for any distribution.

Validated (GB10 2-node here):

• Bit-exact vs torch.topk across 64K–1M × {spread, dense-pivot, >4096-exact-ties} — value multiset identical, no dups, fallback verified.
• Kernel time (single row): 64K 0.038 vs 0.059 ms, 256K 0.078 vs 0.159, 512K 0.128 vs 0.291, 1M 0.229 vs 0.557 ms — 2.43×. The win grows with context (it's read-bound, so halving reads pays off more the longer the row). Only the rare overflow/dense-pivot case is ~12% slower (5 reads: the histogram pass + the streaming fallback).
• Serve: DSv4-Flash at max_model_len=524288, arthur needle-recall at ~434K = 2/2, clean, no oversubscribe/errors — the 2-pass path is correct in the real indexer.

The honest caveat, and where your sweep comes in. This is 2.4× on the top-k. Whether it moves your ~110 s prefill depends on how much of that step the top-k actually is — which is exactly what your prefill-latency-vs-context sweep + a py-spy/nsys sample would show. If the top-k dominates, you should see a real cut; if the sparse-MLA prefill / MoE dominate, less so (and those are the next targets). Would love the before/after on your 4-node/1M rig — pull 9f18be7630 (or I can cut a tag) and re-run the representative window. If it helps end-to-end I'll tag it + write it into the PR notes.

Problem definition: long-context prefill is slow on GB10 (4-node TP=4/EP=4, 1M ctx)

The root: a single large prefill is minutes long — no concurrency involved

One request, ~536K prompt tokens, max_tokens=1, engine otherwise light: prefill (TTFT) = 573 s (~9.5 min). ≈930 tok/s prefill. That's a single request — no parallelism, no scheduler contention — so it isolates the raw kernel cost cleanly. It lines up with your own 2-node number (535K needle = 513 s), so it reproduces across topology.

This is the problem. Everything else is a consequence of it.

Where the ~9.5 min goes

py-spy sampling the worker's MainThread across the prefill (8 samples):

So your histogram_2048 top-k win (2.43× on ~12% of the step) is a real ~7–15% per-prefill cut — but the sparse-MLA attention (~50%) and MoE (~37%) dominate long-context prefill, and look like the higher-leverage targets. The step is also interconnect-latency-bound here (GPUs spin-wait on NCCL over the CX7 fabric), so launch-overhead reductions mostly don't help — it's genuinely kernel/comms time.

Profile caveat: this is py-spy (samples where the Python thread sits — often parked in a cuLaunchKernel), a proxy for GPU time, not a kernel timeline. nsys isn't in the serving container and host-attach needs sudo we don't have on these nodes — so if a real nsys/ncu GPU-kernel breakdown would sharpen the MLA-vs-MoE split, we'll relaunch the serve under nsys in a maintenance window and capture a 1M-prefill trace. Say the word.

Why it surfaces as a "stall with parallel requests"

The several-minute single-prefill is the cost the submitter of a big request pays. It becomes a stall for everyone else under concurrency because max_num_partial_prefills=1 (default) → the engine prefills one request at a time, so N concurrent large requests serialize:

07:31:55  prompt 0.0 tok/s  gen 0.0  Running 2  Waiting 0   <- one prefill monopolizes
07:33:15  prompt 0.0 tok/s  gen 0.3  Running 2  Waiting 4   <- queue backs up ~110s
07:33:55  prompt 11609 tok/s gen 51  Running 3  Waiting 1   <- prefill lands, recovers

So parallel-stall ≈ N × single-prefill-time. Chunked prefill doesn't hide it — the indexer top-k is over the accumulated context, so every late chunk still does ~full-context work.

Fix space

• Kernel (yours, the root): MLA-attention + MoE dominate the ~9.5-min prefill — bigger targets than the top-k for long-context latency. Speeding these helps single-request TTFT and the parallel stall equally.
• Scheduler (ours): max_num_partial_prefills gates the serialization; raising it might overlap prefills, but with the accumulated-context top-k + interconnect-bound step our guess is it'd thrash — an open question we'd value your read on.

This is excellent isolation work — thank you. It matches our read, and it correctly reframes where the top-k sits.

Agreed on the shape. top-k ≈12%, and your ~7–15%-per-prefill math for the 2-pass histogram is right — it was always aimed at the top-k monopolization (the 4-read streaming blowup when a ~1M-context row overflows the register cache), never the prefill floor. MLA + MoE are the floor, and it's interconnect-bound — that lines up with our own 2-node decode teardown (MoE-GEMM + NCCL-AR + dense-GEMM as the top-3), just shifted toward attention on the prefill side.

One caveat before we chase the 50/37 split — py-spy will over-attribute to whichever kernel the NCCL-wait parks in. On an interconnect-bound step the MainThread sits in cuLaunchKernel behind the collective, so a sample lands on whatever kernel is launching when the stall hits — which inflates the two biggest launchers (MLA + MoE) with comms-wait time. Your own caveat, and it's the load-bearing one: if a meaningful slice of the wall-clock is actually NCCL-wait overlapped onto those launches, the true GPU-kernel split — and therefore the target — moves. So yes, please grab the nsys trace. A single-request 536K (or 1M) prefill, max_tokens=1, GPU-kernel timeline is exactly what decides MLA-attn vs MoE-GEMM vs indexer vs comms. That's the gating measurement before either of us commits kernel effort — say the word and we'll take whatever you capture.

One thing already pointing that way: your 4-node 573s is actually a hair slower than the 2-node 513s for the same ~535K — i.e. adding two nodes bought more collective latency than compute relief. For a single fat prefill that's a fingerprint of the comms-bound regime (the fabric, not FLOPs, is the near-term wall), and it says the node count that minimizes single-prefill TTFT may be the smallest one that still holds the KV — the extra nodes are there for 1M capacity, not speed.

One structural note that the trace should confirm: the ~50% "sparse-MLA attention" in prefill is DSv4's MHA-style prefill path — a different kernel from the sparse-decode one the top-k / FI-SM120 work touches (the SM120 sparse backend is decode-only). So it's separately optimizable and, as you say, the higher-leverage target. On the MoE side (~37%), consumer-Blackwell defaults to Marlin W4A16; whether Marlin is the right kernel for these large prefill row-counts (vs a grouped-GEMM path) is an open question, and the natural large-batch MoE lever — DeepGEMM W4A8 — is currently blocked on SM120, so near-term it's Marlin or bust there. Since the step is comms-bound, the one no-rebuild knob worth a cheap A/B is NCCL tuning (algo/proto/channels for the fabric) — it won't touch kernel time but can shave the collective floor.

On max_num_partial_prefills — your thrash instinct is right, for throughput. A single 536K prefill already saturates both the GPU and the fabric, so running two concurrently adds no compute and no bandwidth — it splits the same resource, each prefill runs ~2× slower, and total wall-clock is ~unchanged (worse once the accumulated-context top-k runs for both). p50 for the first submitter degrades, the tail stays flat; serialize-at-1 is actually protecting mean TTFT. What raising it would help is a different goal — fairness for small requests stuck behind a big prefill — and the better lever there is chunked-prefill interleaving + long_prefill_token_threshold to cap the big prefill's per-step budget so small requests slip between its chunks. The catch you already named applies: the indexer top-k is over accumulated context, so every late chunk still pays ~full-context indexer cost → chunking makes the big prefill itself slower while unblocking the small ones (a fairness-vs-big-prefill-throughput trade, tunable via the threshold). It doesn't move the single-prefill floor — only the kernel does, and that (as you noted) helps single-request TTFT and the parallel stall equally.

So: nsys next. A single-request 1M/536K prefill GPU timeline decides whether the effort goes to the prefill-MLA kernel, the MoE GEMM, or the comms floor — grab it whenever the maintenance window opens and we'll target from there.

@epheien Thanks for the report.

I traced this to the strict tool-calling structural-tag grammar for DeepSeek V4 chat-mode outputs. In your log, the rejected token sequence [16, 128822] decodes as .</think>. The prompt in chat mode already ends the thinking section, but the model can occasionally emit an extra </think> before entering normal content/tool-call output. After the rebase, xgrammar's DeepSeek V4 non-reasoning structural tag excluded </think>, so the scheduler treated that otherwise harmless extra close tag as a grammar violation and terminated the request.

I pushed a narrow fix in 531807c34bf73f6ef095f203486e12dadfbdadb2: for DeepSeek V4 non-reasoning strict tool-calling only, allow a stray </think> while still rejecting <think> and leaving the reasoning-mode structural tag unchanged.

Verified with:

• tests/tool_parsers/test_structural_tag_registry.py -q
• tests/tool_parsers/test_deepseekv4_tool_parser.py tests/reasoning/test_deepseekv4_reasoning_parser.py -q
• an xgrammar matcher check for .</think> and .</think><｜DSML｜tool_calls>\n

Please rebuild/redeploy the latest PR head and retry. If it still reproduces, please share the launch args around --tool-call-parser, --reasoning-parser, VLLM_ENFORCE_STRICT_TOOL_CALLING, and MTP/spec decode settings.

@jasl Following up on the strict tool-calling grammar thread above — we hit what looks like a different xgrammar termination on a branch head that already includes 531807c34: with MTP spec decode enabled, named-tool_choice requests with thinking on intermittently 500.

Setup: 2× RTX PRO 6000 workstation (sm_120), TP=2, commit 616a5723ac, DeepSeek-V4-Flash-DSpark.
Launch args: --kv-cache-dtype=fp8 --max-model-len=32768 --enable-auto-tool-choice --tool-call-parser=deepseek_v4 --reasoning-parser=deepseek_v4 --speculative-config={"method": "deepseek_mtp", "num_speculative_tokens": 3}. VLLM_ENFORCE_STRICT_TOOL_CALLING unset.

Symptom: in a 42-request structured-extraction batch (~8k-token prompts, concurrency 10), 7/42 returned 500 Internal server error (~38 s server-side), correlated with bursts of:

/project/cpp/grammar_matcher.cc:612: Warning: The matcher has terminated after accepting the stop token, but is trying to accept new token with id 128825.

Unlike the June report, the rejected ids vary per request (128825, 119061, 112434, 1, …) — it reads like draft tokens still being proposed after the matcher accepted the stop token, not a fixed stray tag.

Triangulation (same workload): MTP off + thinking on → 42/42 OK. MTP on + thinking off (chat_template_kwargs) → 42/42 OK (warnings still logged, requests succeed). MTP on + thinking on → 7/42 → 500. Not reproducible with a short single request — appears load/length-dependent.

Workaround we're shipping: thinking off whenever a function is forced via tool_choice (no exact-match loss on our extraction evals). Happy to provide full logs or test a patch. Otherwise the branch has been excellent on sm_120 — MTP alone takes us 85 → 151 tok/s single-stream.

Disclosure: I'm not an expert in AI hardware or vLLM internals — this report was put together with extensive help from Claude Code. Numbers and logs are from real runs; the root-cause reading may be off, happy to be corrected.

@jedelacarrera I will check tomorrow, thank you for the report!

@jedelacarrera Thank you — this was an excellent report, and your instinct was right on two counts: the workaround is sound, and the grammar_matcher.cc … terminated warning was a red herring (it fires in the passing case too).

Reproduced on a 2-node GB10 (TP=2) with your configuration — MTP + --enable-auto-tool-choice --tool-call-parser deepseek_v4 --reasoning-parser deepseek_v4, thinking on, forced tool_choice, ~8k-token prompts @ concurrency 10: 16/42 → 500. It's an upstream structured-output × spec-decode issue and is independent of hardware (the buggy code lives in the scheduler / structured-output manager), so it applies to your sm120 setup as well.

Root cause (not the terminated matcher): every 500 is scheduler.py … Unexpected: grammar rejected tokens … Terminating request, and the rejected token lists decode to the reasoning tail + </think> — e.g. [270, 4105, 16, 128822] = " the tool.</think>". Under speculative decoding, when the reasoning stream ends inside a single MTP step, the scheduler was advancing the tool-call structural-tag grammar over the pre-</think> reasoning tokens; the grammar correctly rejects them and the request is terminated. That's exactly why it's thinking-on + MTP-on only: thinking-off has no reasoning tokens to mis-feed, and non-MTP never straddles the boundary within a single step.

This is already fixed upstream by #44297 ("Constrain bitmask and trim grammar advance at the reasoning boundary"). I've backported it to this branch (codex/ds4-sm120-min-enable / ds4-sm120-preview-dev, commit fe486bf6). It's a pure-Python change — no kernel rebuild required.

Validated on the same box: the identical 42-request workload → 42/42, 0×500, 0 terminated requests; and #44297's own test suite (tests/v1/spec_decode/test_mtp_structured_output.py, incl. test_bitmask_constrained_when_reasoning_ends_midwindow) passes on our hardware.

Please pull the latest branch head and retry. Your thinking-off-under-forced-tool_choice workaround is a perfectly good stopgap in the meantime. Thanks again — the decoded token ids and the MTP/thinking triangulation in your report made this quick to pin down.

One heads-up: after #44297 you may still see a few non-fatal Failed to advance FSM … Please file an issue. log lines under this workload — that's an expected, tolerated artifact of #44297's bitmask-simulation path (it probes post-reasoning drafts and tolerates rejection, but accept_tokens still logs). Requests succeed; it's cosmetic only.

@jedelacarrera Thank you — this was an excellent report, and your instinct was right on two counts: the workaround is sound, and the grammar_matcher.cc … terminated warning was a red herring (it fires in the passing case too).

Reproduced on a 2-node GB10 (TP=2) with your configuration — MTP + --enable-auto-tool-choice --tool-call-parser deepseek_v4 --reasoning-parser deepseek_v4, thinking on, forced tool_choice, ~8k-token prompts @ concurrency 10: 16/42 → 500. It's an upstream structured-output × spec-decode issue and is independent of hardware (the buggy code lives in the scheduler / structured-output manager), so it applies to your sm120 setup as well.

Root cause (not the terminated matcher): every 500 is scheduler.py … Unexpected: grammar rejected tokens … Terminating request, and the rejected token lists decode to the reasoning tail + </think> — e.g. [270, 4105, 16, 128822] = " the tool.</think>". Under speculative decoding, when the reasoning stream ends inside a single MTP step, the scheduler was advancing the tool-call structural-tag grammar over the pre-</think> reasoning tokens; the grammar correctly rejects them and the request is terminated. That's exactly why it's thinking-on + MTP-on only: thinking-off has no reasoning tokens to mis-feed, and non-MTP never straddles the boundary within a single step.

This is already fixed upstream by #44297 ("Constrain bitmask and trim grammar advance at the reasoning boundary"). I've backported it to this branch (codex/ds4-sm120-min-enable / ds4-sm120-preview-dev, commit fe486bf6). It's a pure-Python change — no kernel rebuild required.

Validated on the same box: the identical 42-request workload → 42/42, 0×500, 0 terminated requests; and #44297's own test suite (tests/v1/spec_decode/test_mtp_structured_output.py, incl. test_bitmask_constrained_when_reasoning_ends_midwindow) passes on our hardware.

Please pull the latest branch head and retry. Your thinking-off-under-forced-tool_choice workaround is a perfectly good stopgap in the meantime. Thanks again — the decoded token ids and the MTP/thinking triangulation in your report made this quick to pin down.

One heads-up: after #44297 you may still see a few non-fatal Failed to advance FSM … Please file an issue. log lines under this workload — that's an expected, tolerated artifact of #44297's bitmask-simulation path (it probes post-reasoning drafts and tolerates rejection, but accept_tokens still logs). Requests succeed; it's cosmetic only.

@jasl Confirmed on sm_120 (2× RTX PRO 6000, TP=2): cherry-picked fe486bf’s three files onto our 616a572 build — the previously-failing workload (thinking on, named tool_choice, MTP, ~8k prompts) now passes clean: [42/42, 0×500 / 10/10 smoke]. One note for others on this path: the other commit on the branch head (30283b9, FlashInfer SM120 default) requires a newer FlashInfer with flashinfer.mla._sparse_mla_sm120 — pulling the full head onto an older image crashes at model load; either update FlashInfer or take #44297’s files only. Thanks for the fast turnaround!

nsys GPU-kernel trace — long-context prefill (the gating measurement)

Build: 9f18be7 (PR #41834), 1M ctx, TP4/EP4, packed FlashInfer SM120 prefill+decode, GB10.
Captures: two single-request prefills, max_tokens=1, unique prompts (no prefix-cache hit),
traced with nsys --trace=cuda wrapping the worker on rank0's GB10. Per-kernel GPU durations
are exact under nsys (overhead is CPU-launch-side); wall-clock inflated but irrelevant to the
% split. GPU stays ≈99.7% busy in both — the fabric is not the wall here.

Headline

Long-context prefill is GPU-COMPUTE-bound (GPU ≈99.7% busy), NOT comms-bound. The dominant
kernel is the lightning-indexer's fp8 MQA-logits — and its share GROWS with context.

GPU-kernel time breakdown (rank0), two context sizes

Combined lightning-indexer (MQA-logits + top-k + compress): 57.5% → 65.0% as context grows.

What this settles

1. Your py-spy caveat was right to flag — but it isn't comms-inflation. The "~50% MLA"
   py-spy bucket is the indexer MQA-logits kernel (52–59%), real GPU compute, not NCCL-wait.
   Comms is only 5–7% of GPU time and the GPU is 99.7% busy → prefill is genuinely
   compute-bound. (Decode stays the interconnect-latency-bound regime; prefill is not.)
2. The scaling is the story. From 200K→536K accumulated context the indexer's share rises
   (52.5%→59.3%) while MoE (14.8%→11.4%) and comms (7.5%→5.4%) shrink. That's the O(context)
   fingerprint: the indexer cost grows ~linearly with accumulated context, everything else is
   ~constant per token. This is exactly why single-request TTFT grows ~linearly with prompt
   length AND why parallel big prefills stall (each late chunk re-scans ≈full context).
3. The actual sparse-MLA attention is only ~9% — the packed FlashInfer SM120 prefill path
   (PREFILL=1) already keeps it cheap. The bottleneck is upstream, in the indexer.
4. Highest-leverage target: _fp8_mqa_logits_kernel. Is there a packed/FlashInfer or fused
   indexer path for prefill, analogous to the SM120 decode indexer? That single kernel is
   59% of a 536K prefill and rising — it moves single-request TTFT and the parallel stall
   together, far more than attention or MoE.
5. NCCL tuning has low prefill upside (comms 5–7%, falling with context). I applied multi-QP
   (NCCL_IB_QPS_PER_CONNECTION=4) and am A/B-ing it, but the trace says the fabric isn't the
   prefill wall — it's the indexer. (Multi-QP may still help decode, which is comms-bound.)

Raw .nsys-rep traces (200K: 13 MB, 536K: 8 MB) available on request.

@wingcomm This is a great trace — a per-kernel nsys split is exactly the measurement that settles the py-spy ambiguity, thank you. It lines up with what we've been seeing and sharpens it in one important way.

On the attribution: our earlier py-spy had bucketed that ~52% as "MLA," but your nsys resolves it correctly — it's the lightning-indexer _fp8_mqa_logits_kernel, and the actual sparse-MLA prefill attention (sparse_mla_prefill_mg_dual_kernel) is only ~9%. That matches our own kernel-level view: the indexer scoring is the O(N²) residual; the packed FlashInfer SM120 prefill path already keeps the attention itself cheap. Your O(context) scaling (52→59% as accumulated ctx grows, MoE/comms shrinking) is the N×N all-pairs signature — attention is ~O(N·d) per token, the indexer is the one that grows.

On your headline ask — a packed/fused prefill indexer, analogous to the SM120 decode indexer: we chased exactly this, and it looks like a hard floor recall-safely, so I'll save you the detour. The decode indexer packs because per step it scores 1×N; prefill is N×N. The sub-quadratic angles we tried all broke recall or were vacuous:

• interval/range-bound pruning → vacuous (≈0% pruned at the recall we need),
• cross-layer top-k reuse → no-go (top-512 overlap only ~38% across layers),
• pooled / block-sparse scoring → recall collapse (0.25–0.39; the pooled keys end up too cos-similar),
• MISA-dagger-style → only ~1.5×, and exact-only.

The kernel itself is already fp8 + tuned and compute-bound (not bandwidth) — fusing the top-k into it is only ~4%. So the realistic prefill wins there are constant-factor, not complexity; the O(N²) is inherent to top-k-selecting over the full context.

On NCCL: agreed, low prefill upside — comms at 5–7% and falling is consistent with what we see (prefill is compute-bound, the fabric isn't the wall). Decode is the other story: that is the interconnect-latency-bound regime for us on 2-node, so I'd be very interested in how your multi-QP (NCCL_IB_QPS_PER_CONNECTION=4) A/B lands on decode specifically — if it moves decode, that's a real no-rebuild lever.

Would love the .nsys-rep traces (both sizes) if you don't mind sharing — the exact per-kernel timings are useful to diff against our own captures. Thanks again, this is high-signal.

Thanks — and thank you for the indexer floor analysis; that saves us a real detour. Agreed
it's the N×N wall: top-k over the full accumulated context is inherent, and if interval-prune
goes vacuous / cross-layer overlap is only ~38% / pooled scoring collapses recall, there's no
sub-quadratic path that keeps recall. We'll treat long-context prefill (and the parallel-prefill
stall) as a fundamental cost and stop hunting a kernel rescue for it. Constant-factor only, as
you say.

On the multi-QP decode A/B you asked for — it's neutral. No measurable movement.

Controlled decode benchmark (short prompt, 128 output tokens, temp 0.7, median of 3 trials/conc)
on the live 4-node/1M build, NCCL_IB_QPS_PER_CONNECTION=4 ON vs the base recipe OFF:

If anything OFF is a hair faster at conc-1, but it's inside trial noise — call it a wash. The
reason lines up with the trace's story: decode ITL here is small-message all-reduce
RTT-latency-bound, and multi-QP is a bandwidth lever, so it has nothing to bite on. Also worth
noting — NCCL is already auto-selecting RING_LL for these collectives (visible in the prefill
trace as ncclDevKernel_AllReduce_Sum_bf16_RING_LL), so the low-latency proto is already in play;
forcing LL128 is largely moot. Net: NCCL knob-tuning isn't the no-rebuild decode lever on this
fabric. I've reverted multi-QP off (kept the config minimal).

The decode numbers for reference: ~50 tok/s/req at conc-1 (ITL ~44 ms, MTP giving ~2 tok/step),
flattening to ~37 tok/s/req by conc-4–8 (ITL ~58 ms) as the per-step collective contends. The
real decode levers that remain look like MTP k↑ (more tokens per fixed-latency collective —
you've got the accept-rate headroom) or lower TP per replica (fewer collectives/token) — not
NCCL flags. Curious whether that matches your 2-node view.

Attaching the full per-kernel GPU-time tables (cuda_gpu_kern_sum, 108 kernels each) for both
sizes as CSV — that's the exact per-kernel timing to diff against your own captures. I left the
raw .nsys-rep out of the public thread since nsys bakes our internal fabric addresses into the
captured command line; if you want the full GUI timeline, say the word and I'll get you a scrubbed
copy or send it your way directly.

`=== multi-QP ON (NCCL_IB_QPS_PER_CONNECTION=4) — decode benchmark, 3 trials/conc ===
run: 2026-07-07 00:38:49
conc=1 ntok=128: per-req 49.52 tok/s | median ITL 43.7 ms | aggregate 47.02 tok/s | n=1
conc=1 ntok=128: per-req 49.37 tok/s | median ITL 44.1 ms | aggregate 47.02 tok/s | n=1
conc=1 ntok=128: per-req 50.27 tok/s | median ITL 44.0 ms | aggregate 47.72 tok/s | n=1
conc=4 ntok=128: per-req 36.92 tok/s | median ITL 58.7 ms | aggregate 80.57 tok/s | n=4
conc=4 ntok=128: per-req 37.43 tok/s | median ITL 58.6 ms | aggregate 79.20 tok/s | n=4
conc=4 ntok=128: per-req 37.87 tok/s | median ITL 58.7 ms | aggregate 81.58 tok/s | n=4
conc=8 ntok=128: per-req 36.89 tok/s | median ITL 58.9 ms | aggregate 93.92 tok/s | n=8
conc=8 ntok=128: per-req 37.41 tok/s | median ITL 58.5 ms | aggregate 95.53 tok/s | n=8
conc=8 ntok=128: per-req 36.68 tok/s | median ITL 58.2 ms | aggregate 93.38 tok/s | n=8

=== multi-QP OFF (base recipe, no QP line) — decode benchmark, 3 trials/conc ===
run: 2026-07-07 00:43:59
conc=1 ntok=128: per-req 52.54 tok/s | median ITL 43.4 ms | aggregate 49.59 tok/s | n=1
conc=1 ntok=128: per-req 49.88 tok/s | median ITL 43.4 ms | aggregate 47.22 tok/s | n=1
conc=1 ntok=128: per-req 50.00 tok/s | median ITL 43.3 ms | aggregate 47.45 tok/s | n=1
conc=4 ntok=128: per-req 37.01 tok/s | median ITL 58.7 ms | aggregate 78.17 tok/s | n=4
conc=4 ntok=128: per-req 38.52 tok/s | median ITL 57.6 ms | aggregate 80.71 tok/s | n=4
conc=4 ntok=128: per-req 36.98 tok/s | median ITL 58.5 ms | aggregate 80.31 tok/s | n=4
conc=8 ntok=128: per-req 37.36 tok/s | median ITL 57.9 ms | aggregate 65.44 tok/s | n=8
conc=8 ntok=128: per-req 38.54 tok/s | median ITL 58.5 ms | aggregate 95.26 tok/s | n=8
conc=8 ntok=128: per-req 35.84 tok/s | median ITL 58.0 ms | aggregate 92.84 tok/s | n=8

=== VERDICT: multi-QP (NCCL_IB_QPS_PER_CONNECTION=4) is NEUTRAL on decode ===
per-req tok/s median ITL (ms)
conc ON OFF ON OFF
1 49.7 50.8 44.0 43.4 (OFF marginally faster; within noise)
4 37.4 37.5 58.7 58.3 (indistinguishable)
8 37.0 37.2 58.5 58.1 (indistinguishable)
Conclusion: no measurable decode benefit. Decode ITL is small-message all-reduce
RTT-latency-bound; multi-QP is a bandwidth lever -> doesn't move it. NCCL already
auto-selects RING_LL (low-latency proto) for these small collectives. Kept multi-QP
OFF (base recipe) as live config. Real decode levers = MTP k-up (more tokens/collective)
or lower TP per replica, NOT NCCL knobs. Prefill separately confirmed compute-bound
(indexer O(N^2)) by nsys, comms 5-7% -> multi-QP no prefill upside either.`

dsv4-9f18be76-prefill-200k-kernels.csv
dsv4-9f18be76-prefill-536k-kernels.csv

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jasl.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork